package cn.lsh.kafka.partitioner;

import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;

import java.util.List;
import java.util.Map;

/**
 * 自定义生产者分区器
 */
public class MyProducerPartitioner implements Partitioner {
	@Override
	public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
		List<PartitionInfo> partitionInfos = cluster.partitionsForTopic(topic);
		//获取可用分区
		List<PartitionInfo> availablePartitionInfos = cluster.availablePartitionsForTopic(topic);
		//获得当前topic的分区数
		Integer pCount = cluster.partitionCountForTopic(topic);
		if (key == null) {
			throw new RuntimeException(topic + "主题key不能为空");
		}
		if (key instanceof Integer) {
			int i = (int) key;
			if (i >= 0 && i < pCount) {
				return i;
			}
		}
		throw new RuntimeException(topic + "主题key必须为小于" + pCount + "的正整数");

	}

	@Override
	public void close() {

	}

	@Override
	public void configure(Map<String, ?> map) {

	}
}
